import argparse
import numpy as np

from tensorflow.keras.models import Sequential
from tensorflow.keras.layers import LSTM, Dense, Reshape, Dropout
from tensorflow.keras.optimizers import Adam
from tensorflow.keras.callbacks import EarlyStopping, ModelCheckpoint

from util.util_log import setup_logging

from multi_beam_trace.util_create_trace import create_scan_traces_theta, \
    create_scan_traces_phi, create_scan_traces_theta_phi

from beam_trace.util_phase_pattern import traces_2_phaseRads, phases_rad_2_bit, phaseBit_2_pattern_point
from beam_trace.util_plot import save_images_2x2


import os
os.environ['KMP_DUPLICATE_LIB_OK'] = 'TRUE'


# ============================================= LSTM相关 =======================================
def prepare_sequences(phaseBits, sequence_length=10):
    """
    将码阵序列转换为LSTM输入格式。
    :param phaseBits: 码阵列表，每个码阵为64x64的二维数组
    :param sequence_length: 历史序列长度
    :return: X (输入序列), y (目标码阵)
    """
    X, y = [], []
    for i in range(len(phaseBits) - sequence_length):
        # 提取历史序列
        seq = phaseBits[i:i + sequence_length]
        # 展平每个码阵为4096维向量
        seq_flat = [arr.flatten() for arr in seq]
        X.append(seq_flat)
        # 目标码阵
        target_flat = phaseBits[i + sequence_length].flatten()
        y.append(target_flat)
    return np.array(X), np.array(y)


def build_lstm_model(input_shape):
    """
    构建LSTM模型。
    :param input_shape: 输入形状 (sequence_length, 4096)
    :return: LSTM模型
    """
    model = Sequential([
        LSTM(256, input_shape=input_shape, return_sequences=False),
        Dropout(0.3),
        Dense(4096, activation='sigmoid'),
    ])
    model.compile(optimizer=Adam(learning_rate=0.001),
                  loss='binary_crossentropy',
                  metrics=['accuracy'])
    return model


def train_model(model, X_train, y_train, path_model_checkpoint, epochs=100, batch_size=32):
    """
    训练模型。
    :param model: LSTM模型
    :param X_train: 训练输入
    :param y_train: 训练目标
    :param epochs: 训练轮数
    :param batch_size: 批次大小
    :return: 训练历史
    """
    callbacks = [
        EarlyStopping(patience=10, monitor='val_loss', restore_best_weights=True),
        ModelCheckpoint(path_model_checkpoint, save_best_only=True)
    ]
    history = model.fit(
        X_train, y_train,
        epochs=epochs,
        batch_size=batch_size,
        validation_split=0.2,
        callbacks=callbacks,
        verbose=1
    )
    return history


# ============================================= 结果相关 =======================================
def save_results(base_dir, y_real_arrays, y_out_arrays, bit_num):
    dir_compare = base_dir + "/compare/"
    total_correct = 0
    total_elements = 0
    accuracies = []
    for idx, (y_real, y_out) in enumerate(zip(y_real_arrays, y_out_arrays)):
        # 计算准确率
        correct = np.sum(y_real == y_out)
        total_correct += correct
        total_elements += y_real.size

        accuracy = correct / y_real.size
        accuracies.append(accuracy)

        if idx % 10 == 0:
            # 计算方向图, 指向, PSLL
            pattern_real, point_real, psll_real = phaseBit_2_pattern_point(y_real, bit_num)
            pattern_out, point_out, psll_out = phaseBit_2_pattern_point(y_out, bit_num)

            # 保存比较图片
            # plot_images_2x2(y_real, y_out, pattern_real, pattern_out,
            #                 f"Accuracy = {accuracy:.2%}, "
            #                 f"phaseBit(real):[{point_real[1]}]-[{psll_real[0]:.2f},{psll_real[1]}], "
            #                 f"phaseBit(out):[{point_out[1]}]-[{psll_out[0]:.2f},{psll_out[1]}]")
            save_images_2x2(dir_compare + str(idx) + ".jpg",
                            y_real, f"phaseBit(real):[{point_real[1]}]-[{psll_real[0]:.2f},{psll_real[1]}]",
                            y_out, f"phaseBit(out):[{point_out[1]}]-[{psll_out[0]:.2f},{psll_out[1]}]",
                            pattern_real, "pattern(real)", pattern_out, "pattern(out)",
                            f"Accuracy = {accuracy:.2%}")

        # 保存.csv
        # np.savetxt(os.path.join(base_dir, f'y_real_{idx}.csv'), y_real, delimiter=',')
        # np.savetxt(os.path.join(base_dir, f'y_out_{idx}.csv'), y_out, delimiter=',')

    # 计算当前数据集的统计信息
    overall_accuracy = total_correct / total_elements
    logger.info(f"\n总体准确率: {overall_accuracy:.2%}")
    logger.info(f"准确率统计:")

    stats = {
        'max': np.max(accuracies),
        'min': np.min(accuracies),
        'mean': np.mean(accuracies),
        'std': np.std(accuracies)
    }

    for metric, value in stats.items():
        logger.info(f"{metric.capitalize()}: {value:.4f}")

    return stats


def save_statistics(statistics, file_path):
    with open(file_path, 'w') as f:
        for key, value in statistics.items():
            f.write(f'{key.capitalize()}: {value}\n')


# ============================================= 主流程相关 =======================================
def get_traces_phaseBits_by_theta(bit_num, theta_start, theta_end, phi_start, phi_end):
    """ 获取波束指向和对应码阵 """
    logger.info(f"theta_start={theta_start}, theta_end={theta_end}, phi_start={phi_start}, phi_end={phi_end}")
    traces = []
    phaseRads = []
    phaseBits = []
    for phi in range(phi_start, phi_end, 1):
        # 生成轨迹指向
        traces_idx = create_scan_traces_theta(theta_start, theta_end, phi)  # theta方向
        # 生成码阵
        phaseRads_idx = traces_2_phaseRads(traces_idx)
        # 码阵bit化
        phaseBits_idx = phases_rad_2_bit(phaseRads_idx, bit_num)
        # 记录
        traces.extend(traces_idx)
        phaseRads.extend(phaseRads_idx)
        phaseBits.extend(phaseBits_idx)
    logger.info(f"len of traces: {len(traces)}")
    return traces, phaseRads, phaseBits

def get_traces_phaseBits_by_phi(bit_num, theta_start, theta_end, phi_start, phi_end):
    """ 获取波束指向和对应码阵 """
    logger.info(f"theta_start={theta_start}, theta_end={theta_end}, phi_start={phi_start}, phi_end={phi_end}")
    traces = []
    phaseRads = []
    phaseBits = []
    for theta in range(theta_start, theta_end, 1):
        # 生成轨迹指向
        traces_idx = create_scan_traces_phi(theta=theta, phi_start=phi_start, phi_end=phi_end)
        # 生成码阵
        phaseRads_idx = traces_2_phaseRads(traces_idx)
        # 码阵bit化
        phaseBits_idx = phases_rad_2_bit(phaseRads_idx, bit_num)
        # 记录
        traces.extend(traces_idx)
        phaseRads.extend(phaseRads_idx)
        phaseBits.extend(phaseBits_idx)
    logger.info(f"len of traces: {len(traces)}")
    return traces, phaseRads, phaseBits


def main(path_dir, bit_num, mode, theta_start, theta_end, phi_start, phi_end):
    # 获取数据
    if mode == 1:
        traces, phaseRads, phaseBits \
            = get_traces_phaseBits_by_theta(bit_num, theta_start, theta_end, phi_start, phi_end)
    elif mode == 2:
        traces, phaseRads, phaseBits \
            = get_traces_phaseBits_by_phi(bit_num, theta_start, theta_end, phi_start, phi_end)
    else:
        traces, phaseRads, phaseBits \
            = get_traces_phaseBits_by_phi(bit_num, theta_start, theta_end, phi_start, phi_end)

    # 数据预处理
    sequence_length = 10  # 历史序列长度
    X, y = prepare_sequences(phaseBits, sequence_length)

    # 划分训练集和测试集
    split = int(0.8 * len(X))
    X_train, X_test = X[:split], X[split:]
    y_train, y_test = y[:split], y[split:]

    # 构建模型
    input_shape = (sequence_length, 64 * 64)
    model = build_lstm_model(input_shape)
    model.summary()

    # 训练模型
    history = train_model(model, X_train, y_train, path_dir+'/best_lstm_model.h5', epochs=100, batch_size=32)

    # 在训练集和测试集上都进行预测
    datasets = {
        'train': {'X': X_train, 'y': y_train},
        'test': {'X': X_test, 'y': y_test}
    }
    all_accuracies = {}

    for name, dataset in datasets.items():
        logger.info(f"\n===== 在{name}上进行预测 =====")
        X = dataset['X']
        y = dataset['y']
        # 获取模型预测结果
        y_out = model.predict(X)
        # 二值化预测结果
        y_out_binary = (y_out > 0.5).astype(np.int32)
        # 将LSTM预测结果还原为原始形状
        y_2d = [arr.reshape(64, 64) for arr in y]
        y_out_2d = [arr.reshape(64, 64) for arr in y_out_binary]

        # 保存结果并获取统计信息
        stats = save_results(path_dir + f'/dataset_{name}', y_2d, y_out_2d, bit_num)
        all_accuracies[name] = stats

    # 保存统计信息
    save_statistics({
        'train_max': all_accuracies['train']['max'],
        'train_min': all_accuracies['train']['min'],
        'train_mean': all_accuracies['train']['mean'],
        'train_std': all_accuracies['train']['std'],
        'test_max': all_accuracies['test']['max'],
        'test_min': all_accuracies['test']['min'],
        'test_mean': all_accuracies['test']['mean'],
        'test_std': all_accuracies['test']['std']
    }, path_dir + "/result.txt")

    # 保存模型
    model.save(path_dir+'/lstm_phaseBits_predictor.h5')


if __name__ == "__main__":
    # 创建参数解析器
    parser = argparse.ArgumentParser(description="Process some parameters.")
    parser.add_argument("--base_path", type=str,
                        default="../files/dissertation/chapter_3/tf-ae1d-bit-cnn2d-(1,90)",
                        help="Base directory path (default: ../files/dissertation/chapter_3/tf-ae1d-bit-cnn2d-(1,90))")
    parser.add_argument("--bit_num", type=int, default=1, help="Number of bits (default: 1)")
    parser.add_argument("--mode", type=int, default=2, help="1: by theta, 2: by phi (default: 2)")
    parser.add_argument("--theta_start", type=int, default=1, help="theta_start (default: 1)")
    parser.add_argument("--theta_end", type=int, default=60, help="theta_end (default: 60)")
    parser.add_argument("--phi_start", type=int, default=0, help="phi_start (default: 0)")
    parser.add_argument("--phi_end", type=int, default=360, help="phi_end (default: 360)")

    args = parser.parse_args()

    base_path = args.base_path
    mode = args.mode
    bit_num = args.bit_num
    theta_start, theta_end, phi_start, phi_end = args.theta_start, args.theta_end, args.phi_start, args.phi_end

    # 初始化日志
    logger = setup_logging(base_path + "/trace.txt")
    # 示例日志记录
    logger.info(f"Starting execution with base_path: {base_path}")
    logger.info(f"Using bit_num: {bit_num}")
    logger.info(f"theta_start={theta_start}, theta_end={theta_end}, phi_start={phi_start}, phi_end={phi_end}")

    main(path_dir=base_path, bit_num=bit_num, mode=mode,
         theta_start=theta_start, theta_end=theta_end, phi_start=phi_start, phi_end=phi_end)
